//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package parsecmd

import (
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"github.com/jhump/protoreflect/desc"
	"github.com/jhump/protoreflect/desc/protoparse"
	"github.com/spf13/cobra"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/descriptorpb"
)

var rootPath = ""
var importPathList []string
var outPath = ""

var ParseCommand = &cobra.Command{
	Use: "parse",
	RunE: func(cmd *cobra.Command, args []string) error {
		return runParseCommand()
	},
}

func runParseCommand() error {
	fileNames := []string{}
	filepath.WalkDir(rootPath, func(path string, _ fs.DirEntry, fsErr error) error {
		if fsErr != nil { //skip error
			return nil
		}
		if strings.HasSuffix(path, ".proto") {
			if !strings.Contains(path, "google") { //忽略google相关标准库
				fileNames = append(fileNames, path)
			}
		}
		return nil
	})
	tmpImportList := []string{}
	tmpImportList = append(tmpImportList, rootPath)
	tmpImportList = append(tmpImportList, importPathList...)
	fileNames, err := protoparse.ResolveFilenames(tmpImportList, fileNames...)
	if err != nil {
		return err
	}
	p := protoparse.Parser{
		ImportPaths:           tmpImportList,
		InferImportPaths:      len(tmpImportList) == 0,
		IncludeSourceCodeInfo: true,
	}
	fds, err := p.ParseFiles(fileNames...)
	if err != nil {
		return err
	}
	tmpMap := map[string]*desc.FileDescriptor{}
	for _, fd := range fds {
		if err := addFile(fd, tmpMap); err != nil {
			return err
		}
	}
	expandedFiles := map[string]struct{}{}
	allFilesSlice := []*descriptorpb.FileDescriptorProto{}
	for name := range tmpMap {
		allFilesSlice = addFilesToSet(allFilesSlice, expandedFiles, tmpMap[name])
	}
	data, err := proto.Marshal(&descriptorpb.FileDescriptorSet{File: allFilesSlice})
	if err != nil {
		return err
	}
	err = os.WriteFile(outPath, data, 0600)
	if err != nil {
		return err
	}
	fmt.Println("{}")
	return nil
}

func addFilesToSet(allFiles []*descriptorpb.FileDescriptorProto, expanded map[string]struct{}, fd *desc.FileDescriptor) []*descriptorpb.FileDescriptorProto {
	if _, ok := expanded[fd.GetName()]; ok {
		// already seen this one
		return allFiles
	}
	expanded[fd.GetName()] = struct{}{}
	// add all dependencies first
	for _, dep := range fd.GetDependencies() {
		allFiles = addFilesToSet(allFiles, expanded, dep)
	}
	return append(allFiles, fd.AsFileDescriptorProto())
}

func addFile(fd *desc.FileDescriptor, fds map[string]*desc.FileDescriptor) error {
	name := fd.GetName()
	if existing, ok := fds[name]; ok {
		// already added this file
		if existing != fd {
			// doh! duplicate files provided
			return fmt.Errorf("given files include multiple copies of %q", name)
		}
		return nil
	}
	fds[name] = fd
	for _, dep := range fd.GetDependencies() {
		if err := addFile(dep, fds); err != nil {
			return err
		}
	}
	return nil
}

func init() {
	ParseCommand.Flags().StringVar(&rootPath, "rootPath", "", "grpc proto file root path")
	ParseCommand.Flags().StringArrayVar(&importPathList, "importPath", []string{}, "grpc import path")
	ParseCommand.Flags().StringVar(&outPath, "outPath", "", "grpc proto parse result save path")
}
